#include <gtest/gtest.h>
#include "Thread/ThreadPool.h"
#include <chrono>
#include <atomic>
#include <vector>
#include <future>

using namespace toolkit;
using namespace std::chrono_literals;

// 测试线程池基本功能
TEST(ThreadPoolTest, BasicFunctionality) {
    ThreadPool pool(4); // 创建4个线程的线程池
    
    // 初始负载应为0
    EXPECT_EQ(pool.getLoad(), 0);
    EXPECT_EQ(pool.size(), 4);
    
    std::atomic<int> counter{0};
    std::vector<TaskCancelable::Ptr> tasks;
    
    // 提交10个任务
    for (int i = 0; i < 10; i++) {
        auto task = pool.async([&counter, i]() {
            std::this_thread::sleep_for(50ms);
            counter++;
        });
        tasks.push_back(task);
    }
    
    // 等待所有任务完成
    std::this_thread::sleep_for(600ms);
    
    // 所有任务应该已完成
    EXPECT_EQ(counter, 10);
    
    // 关闭线程池
    pool.shutdown();
}

// 测试任务取消功能
TEST(ThreadPoolTest, TaskCancellation) {
    ThreadPool pool(2);
    
    std::atomic<int> counter{0};
    std::atomic<bool> long_task_running{false};
    
    // 提交一个长时间运行的任务
    auto long_task = pool.async([&]() {
        long_task_running = true;
        std::this_thread::sleep_for(500ms);
        counter++;
    });
    
    // 等待长任务开始运行
    while (!long_task_running) {
        std::this_thread::sleep_for(10ms);
    }
    
    // 提交多个短任务，这些任务会被取消
    std::vector<TaskCancelable::Ptr> tasks;
    for (int i = 0; i < 5; i++) {
        auto task = pool.async([&counter]() {
            std::this_thread::sleep_for(100ms);
            counter++;
        });
        tasks.push_back(task);
    }
    
    // 取消所有短任务
    for (auto& task : tasks) {
        task->cancel();
    }
    
    // 等待足够长的时间让所有未取消的任务完成
    std::this_thread::sleep_for(600ms);
    
    // 只有长任务应该完成
    EXPECT_EQ(counter, 1);
    
    pool.shutdown();
}

// 测试优先任务功能
TEST(ThreadPoolTest, PriorityTask) {
    ThreadPool pool(1); // 单线程池，确保任务按顺序执行
    
    std::vector<int> execution_order;
    std::mutex mutex;
    
    // 提交5个普通任务
    for (int i = 0; i < 5; i++) {
        pool.async([i, &execution_order, &mutex]() {
            std::this_thread::sleep_for(50ms);
            std::lock_guard<std::mutex> lock(mutex);
            execution_order.push_back(i);
        });
    }
    
    // 提交一个优先任务
    pool.async_first([&execution_order, &mutex]() {
        std::lock_guard<std::mutex> lock(mutex);
        execution_order.push_back(999); // 特殊标记
    });
    
    // 等待所有任务完成
    std::this_thread::sleep_for(400ms);
    
    // 检查执行顺序
    std::lock_guard<std::mutex> lock(mutex);
    EXPECT_EQ(execution_order.size(), 6);
    
    // 优先任务应该在前面执行
    // 注意：由于第一个任务可能已经开始执行，优先任务可能不是第一个
    bool priority_found = false;
    for (size_t i = 0; i < execution_order.size(); i++) {
        if (execution_order[i] == 999) {
            priority_found = true;
            EXPECT_LT(i, 2); // 应该在前两个位置
            break;
        }
    }
    EXPECT_TRUE(priority_found);
    
    pool.shutdown();
}

// 测试同步执行功能
TEST(ThreadPoolTest, SyncExecution) {
    // 创建一个没有线程的线程池，所有任务都会同步执行
    ThreadPool pool(0);
    
    std::atomic<bool> task_executed{false};
    
    // 提交任务
    pool.async([&task_executed]() {
        task_executed = true;
    }, true);
    
    // 任务应该已经同步执行完成
    EXPECT_TRUE(task_executed);
    
    pool.shutdown();
}

// 测试负载计数
TEST(ThreadPoolTest, LoadCounting) {
    ThreadPool pool(2);
    
    std::atomic<bool> task1_running{false};
    std::atomic<bool> task2_running{false};
    std::atomic<bool> can_finish{false};
    
    // 提交两个长时间运行的任务
    auto task1 = pool.async([&]() {
        task1_running = true;
        while (!can_finish) {
            std::this_thread::sleep_for(10ms);
        }
    });
    
    auto task2 = pool.async([&]() {
        task2_running = true;
        while (!can_finish) {
            std::this_thread::sleep_for(10ms);
        }
    });
    
    // 等待任务开始执行
    while (!task1_running || !task2_running) {
        std::this_thread::sleep_for(10ms);
    }
    
    // 检查负载
    EXPECT_EQ(pool.getLoad(), 2);
    
    // 允许任务完成
    can_finish = true;
    
    // 等待任务完成
    std::this_thread::sleep_for(100ms);
    
    // 负载应该回到0
    EXPECT_EQ(pool.getLoad(), 0);
    
    pool.shutdown();
}

// 测试大量任务
TEST(ThreadPoolTest, ManyTasks) {
    ThreadPool pool(4);
    
    std::atomic<int> counter{0};
    constexpr int TASK_COUNT = 1000;
    
    // 提交大量任务
    for (int i = 0; i < TASK_COUNT; i++) {
        pool.async([&counter]() {
            counter++;
        });
    }
    
    // 等待所有任务完成
    std::this_thread::sleep_for(1s);
    
    // 所有任务应该已完成
    EXPECT_EQ(counter, TASK_COUNT);
    
    pool.shutdown();
}

int main(int argc, char **argv) {
    ::testing::InitGoogleTest(&argc, argv);
    return RUN_ALL_TESTS();
}